package com.plugin.uniapp.video;

import android.util.Log;

import java.io.ByteArrayOutputStream;
import java.util.Arrays;

public class NalUnitParser {
    public static int[] parseSps(byte[] sps) {
        try {
            BitReader br = new BitReader(sps);

            br.readBits(8); // nal_header

            int profile_idc = br.readBits(8);
            br.readBits(8); // constraint flags + reserved bits
            br.readBits(8); // level_idc
            br.readUE();    // seq_parameter_set_id

            int chroma_format_idc = 1;
            int separate_color_plane_flag = 0;

            if (profile_idc == 100 || profile_idc == 110 || profile_idc == 122 ||
                    profile_idc == 244 || profile_idc == 44 || profile_idc == 83 ||
                    profile_idc == 86 || profile_idc == 118 || profile_idc == 128 ||
                    profile_idc == 138 || profile_idc == 139 || profile_idc == 134) {
                chroma_format_idc = br.readUE();
                if (chroma_format_idc == 3) {
                    separate_color_plane_flag = br.readBit();
                }
                br.readUE(); // bit_depth_luma_minus8
                br.readUE(); // bit_depth_chroma_minus8
                br.readBit(); // qpprime_y_zero_transform_bypass_flag
                boolean seq_scaling_matrix_present_flag = br.readBit() == 1;
                if (seq_scaling_matrix_present_flag) {
                    int scalingListCount = (chroma_format_idc != 3) ? 8 : 12;
                    for (int i = 0; i < scalingListCount; i++) {
                        boolean scalingListPresent = br.readBit() == 1;
                        if (scalingListPresent) {
                            // 忽略 scaling list
                            // 实际实现可省略，除非你要做压缩优化
                        }
                    }
                }
            }

            br.readUE(); // log2_max_frame_num_minus4
            int pic_order_cnt_type = br.readUE();
            if (pic_order_cnt_type == 0) {
                br.readUE(); // log2_max_pic_order_cnt_lsb_minus4
            } else if (pic_order_cnt_type == 1) {
                br.readBit(); // delta_pic_order_always_zero_flag
                br.readSE();  // offset_for_non_ref_pic
                br.readSE();  // offset_for_top_to_bottom_field
                int num_ref_frames_in_pic_order_cnt_cycle = br.readUE();
                for (int i = 0; i < num_ref_frames_in_pic_order_cnt_cycle; i++) {
                    br.readSE(); // offset_for_ref_frame[i]
                }
            }

            br.readUE(); // max_num_ref_frames
            br.readBit(); // gaps_in_frame_num_value_allowed_flag

            int pic_width_in_mbs_minus1 = br.readUE();
            int pic_height_in_map_units_minus1 = br.readUE();

            boolean frame_mbs_only_flag = br.readBit() == 1;
            if (!frame_mbs_only_flag) {
                br.readBit(); // mb_adaptive_frame_field_flag
            }

            br.readBit(); // direct_8x8_inference_flag
            boolean frame_cropping_flag = br.readBit() == 1;

            int frame_crop_left_offset = 0;
            int frame_crop_right_offset = 0;
            int frame_crop_top_offset = 0;
            int frame_crop_bottom_offset = 0;

            if (frame_cropping_flag) {
                frame_crop_left_offset = br.readUE();
                frame_crop_right_offset = br.readUE();
                frame_crop_top_offset = br.readUE();
                frame_crop_bottom_offset = br.readUE();
            }

            int width_in_mbs = pic_width_in_mbs_minus1 + 1;
            int height_in_map_units = pic_height_in_map_units_minus1 + 1;
            int frame_height = height_in_map_units * (frame_mbs_only_flag ? 1 : 2);

            // 计算 crop unit size
            int crop_unit_x = 1;
            int crop_unit_y = 2; // 默认是 4:2:0

            if (chroma_format_idc == 0) {
                crop_unit_x = 1;
                crop_unit_y = 2;
            } else if (chroma_format_idc == 1) {
                crop_unit_x = 2;
                crop_unit_y = 2;
            } else if (chroma_format_idc == 2) {
                crop_unit_x = 2;
                crop_unit_y = 1;
            } else if (chroma_format_idc == 3) {
                crop_unit_x = 1;
                crop_unit_y = 1;
            }

            if (separate_color_plane_flag == 1) {
                crop_unit_y = 1;
            }

            int width = width_in_mbs * 16 - crop_unit_x * (frame_crop_left_offset + frame_crop_right_offset);
            int height = frame_height * 16 - crop_unit_y * (frame_crop_top_offset + frame_crop_bottom_offset);

            return new int[]{width, height};

        } catch (Exception e) {
            Log.e("NALParse", "解析 SPS 失败：" + e.getMessage());
            return null;
        }
    }

    public static byte[] cleanSps(byte[] rawSps) {
        // 去除起始码
        int offset = 0;
        while (offset < rawSps.length - 4) {
            if (rawSps[offset] == 0x00 && rawSps[offset + 1] == 0x00) {
                if (rawSps[offset + 2] == 0x00 && rawSps[offset + 3] == 0x01) {
                    offset += 4;
                    break;
                } else if (rawSps[offset + 2] == 0x01) {
                    offset += 3;
                    break;
                }
            }
            offset++;
        }

        byte[] spsWithoutStartCode = Arrays.copyOfRange(rawSps, offset, rawSps.length);

        // 去除转义字节 0x03
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        for (int i = 0; i < spsWithoutStartCode.length; i++) {
            if (i > 1 && spsWithoutStartCode[i - 2] == 0x00 && spsWithoutStartCode[i - 1] == 0x00 && spsWithoutStartCode[i] == 0x03) {
                // skip emulation prevention byte
                continue;
            }
            out.write(spsWithoutStartCode[i]);
        }

        return out.toByteArray();
    }
}